import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'

import xml.etree.ElementTree as ET

import tensorflow as tf
from PIL import Image
from tqdm import tqdm

from utils.utils import get_classes
from utils.utils_map import get_coco_map, get_map
from yolo import YOLO
import test_img_call as tt


def StepMode1(image_ids, class_names, args):
    confidence = args['confidence']
    nms_iou = args['nms_iou']
    VOCdevkit_path = args['VOCdevkit_path']
    image_all_path = args['image_all_path']
    map_out_path = args['map_out_path']

    print("Load model.")
    yolo = YOLO(confidence=confidence, nms_iou=nms_iou)
    print("Load model done.")

    print("Get predict result.")
    for image_id in tqdm(image_ids):
        image_path = os.path.join(VOCdevkit_path, image_all_path, image_id + ".jpg")
        # print(image_path)
        # image = Image.open(image_path)

        # yolo.get_map_txt(image_id, image, class_names, map_out_path)

        stat, imgcv, detects, fileSaveDrawBox = tt.test_from_web(fileImageName=image_path,
                                                                 detect_result_path='',
                                                                 isSaveFiles=False,
                                                                 isSaveImg=True,##保存圖片
                                                                 showTimerString=False)

        f = open(os.path.join(map_out_path, "detection-results/" + image_id + ".txt"), "w")
        for detect in detects:
            c = detect[0]
            score = detect[1]
            left = detect[2]
            top = detect[3]
            right = detect[4]
            bottom = detect[5]

            predicted_class = class_names[int(c)]

            if predicted_class not in class_names:
                continue

            f.write("%s %s %s %s %s %s\n" % (
                predicted_class, score, str(int(left)), str(int(top)), str(int(right)), str(int(bottom))))
        f.close()

    print("Get predict result done.")


def StepMode2(image_ids, class_names, args):
    anno_path = args['anno_path']
    VOCdevkit_path = args['VOCdevkit_path']
    map_out_path = args['map_out_path']

    print("Get ground truth result.")

    for image_id in tqdm(image_ids):
        with open(os.path.join(map_out_path, "ground-truth/" + image_id + ".txt"), "w") as new_f:
            root = ET.parse(os.path.join(VOCdevkit_path, anno_path + image_id + ".xml")).getroot()
            # print(root)
            for obj in root.findall('object'):
                difficult_flag = False
                if obj.find('difficult') != None:
                    difficult = obj.find('difficult').text
                    if int(difficult) == 1:
                        difficult_flag = True
                obj_name = obj.find('name').text
                if obj_name not in class_names:
                    continue
                bndbox = obj.find('bndbox')
                left = bndbox.find('xmin').text
                top = bndbox.find('ymin').text
                right = bndbox.find('xmax').text
                bottom = bndbox.find('ymax').text

                if difficult_flag:
                    new_f.write("%s %s %s %s %s difficult\n" % (obj_name, left, top, right, bottom))
                else:
                    new_f.write("%s %s %s %s %s\n" % (obj_name, left, top, right, bottom))
    print("Get ground truth result done.")
